import { cn } from '@baml/ui/lib/utils';
import { atom, useAtomValue } from 'jotai';
import { useMemo, useState, useEffect } from 'react';
import React from 'react';
import { displaySettingsAtom } from '../preview-toolbar';
import { getHighlightedParts } from './highlight-utils';
import { TokenEncoderCache } from './render-tokens';
import type { Tiktoken } from 'js-tiktoken/lite';

export const showTokensAtom = atom(
  (get) => get(displaySettingsAtom).showTokens,
);

const HighlightedText: React.FC<{
  text: string;
  highlightChunks: string[];
}> = ({ text, highlightChunks }) => {
  const parts = getHighlightedParts(text, highlightChunks);

  return (
    <>
      {parts.map((part, i) =>
        part.highlight ? (
          <mark
            key={`${i}-${part.highlight}-${part.text.length}`}
            className={cn(
              'inline whitespace-pre-wrap break-words rounded px-1 py-0.5 font-normal text-xs text-primary-foreground',
              part.text.trim() === '' ? 'bg-chart-5/30' : 'bg-chart-1/40',
            )}
          >
            {part.text}
          </mark>
        ) : (
          <React.Fragment key={`${i}-normal-${part.text.length}`}>
            {part.text}
          </React.Fragment>
        ),
      )}
    </>
  );
};

export const RenderPromptPart: React.FC<{
  text: string;
  highlightChunks?: string[];
  model?: string;
  provider?: string;
}> = ({ text, highlightChunks = [], model, provider }) => {
  const showTokens = useAtomValue(showTokensAtom);
  const [encoder, setEncoder] = useState<Tiktoken | null>(null);
  const [isLoading, setIsLoading] = useState(false);

  // Load encoder asynchronously when needed
  useEffect(() => {
    if (!showTokens) {
      setEncoder(null);
      return;
    }

    const encodingName = TokenEncoderCache.getEncodingNameForModel(
      'baml-openai-chat',
      'gpt-4o',
    );
    if (!encodingName) return;

    // Check if already cached
    const cached = TokenEncoderCache.INSTANCE.getEncoder(encodingName);
    if (cached) {
      setEncoder(cached);
      return;
    }

    // Load asynchronously
    setIsLoading(true);
    TokenEncoderCache.INSTANCE.getEncoderAsync(encodingName)
      .then(enc => {
        setEncoder(enc);
        setIsLoading(false);
      })
      .catch(err => {
        console.error('Failed to load tokenizer:', err);
        setIsLoading(false);
      });
  }, [showTokens, model, provider]);

  const tokenizer = useMemo(() => {
    if (!showTokens || !encoder) return undefined;
    return { enc: encoder, tokens: encoder.encode(text) };
  }, [text, showTokens, encoder]);

  // Only compute highlighted text if we're not tokenizing
  const renderContent = useMemo(() => {
    // Show loading state while encoder is being loaded
    if (showTokens && isLoading) {
      return <span className="text-muted-foreground">Loading tokenizer...</span>;
    }

    if (tokenizer) {
      const tokenized = Array.from(tokenizer.tokens).map((token) =>
        tokenizer.enc.decode([token]),
      );
      return (
        <>
          {tokenized.map((token, i) => (
            <span
              key={`${i}-token-${token.length}-${token.charCodeAt(0) || 0}`}
              className={cn(
                'text-white',
                // Uncomment and use these classes if you want to color-code tokens
                [
                  'bg-fuchsia-800',
                  'bg-emerald-700',
                  'bg-yellow-600',
                  'bg-red-700',
                  'bg-cyan-700',
                ][i % 5],
              )}
            >
              {token}
            </span>
          ))}
        </>
      );
    }

    // Only do highlighting if we're not tokenizing
    return <HighlightedText text={text} highlightChunks={highlightChunks} />;
  }, [text, highlightChunks, tokenizer, showTokens, isLoading]);

  return (
    <div className="flex flex-col min-w-0">
      <div className="px-3 pb-3 pt-0 bg-accent group max-h-[600px] overflow-y-auto overflow-x-hidden min-w-0">
        <pre
          className={cn(
            'whitespace-pre-wrap text-xs leading-relaxed transition-all text-primary-foreground min-w-0',
          )}
        >
          {renderContent}
        </pre>
      </div>
    </div>
  );
};
